import cv2
import numpy as np
from sklearn.cluster import KMeans
import gradio as gr
from PIL import Image


def extract_color_palette(image, k=5):
    # 将图像转换为 RGB 格式
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # 将图片转换为一维数组
    reshaped_image = image.reshape(-1, 3)

    # 使用 KMeans 算法进行颜色聚类
    kmeans = KMeans(n_clusters=k, random_state=0).fit(reshaped_image)

    # 获取聚类中心，即主色调
    palette = kmeans.cluster_centers_.astype(int)

    # 检查颜色有效性
    valid_palette = []
    for color in palette:
        if len(color) == 3 and all(0 <= c < 256 for c in color):
            valid_palette.append(tuple(color))
        else:
            print(f"Invalid color detected: {color}")

    # 创建调色板图像
    palette_image = Image.new("RGB", (100 * len(valid_palette), 100))
    for idx, color in enumerate(valid_palette):
        # 确保颜色是有效的 RGB 值
        if len(color) == 3:
            color_image = Image.new("RGB", (100, 100), color)
            palette_image.paste(color_image, (idx * 100, 0))
    # 生成十六进制调色板字符串
    hex_palette = ['#{:02x}{:02x}{:02x}'.format(color[0], color[1], color[2]) for color in valid_palette]

    return palette_image, "\n".join(hex_palette)


# Gradio 界面
with gr.Blocks() as demo:
    gr.Markdown("# KMeans 调色板提取器")
    with gr.Row():
        image_input = gr.Image(type="numpy", label="上传图像", shape=(512, 512))
        output_palette = gr.Image(label="调色板", type="pil")
        output_hex = gr.Textbox(label="色彩十六进制码", interactive=False)

    clear_button = gr.Button("Clear")
    submit_button = gr.Button("Submit")


    # 按钮事件
    def process_image(image):
        palette_image, hex_palette = extract_color_palette(image)
        return palette_image, hex_palette


    submit_button.click(process_image, inputs=image_input, outputs=[output_palette, output_hex])
    clear_button.click(lambda: (None, ""), inputs=[], outputs=[output_palette, output_hex])

demo.launch()